#-*- coding: utf8 -*-
import itertools
import numpy as np
import operator as op

__all__ = [ 'expand_kwds', 'iter_kwds' ]

def expand_kwds(kwds, keys_equivalence={}, valid_keys=None):
    '''
    In some functions, it is desirable to have short versions of keyword 
    arguments for convenience. This function takes a dictionary of keywords and
    expand all the short versions using the dictionary 'keys_equivalence'.
    
    Additionally, it can also test if all keywords in 'kwds' belong to a given
    set 'valid_keys'.
    
    Example
    -------
    
    # Prepare the keys_equivalence and valid_keys parameters
    >>> equiv = { 'c': 'color', 'lw': 'linewidth' }
    >>> keys = set([ 'color', 'linewidth', 'disp' ])

    # Expand short versions
    >>> expand_kwds({'c': 'k'}, equiv, keys)
    {'color': 'k'}
    
    # Prevents invalid keywords from being passed
    >>> expand_kwds({'c': 'k', 'x': 1}, equiv, keys)
    Traceback (most recent call last):
    ...
    TypeError: Invalid keyword argument, 'x'
    '''

    for (k, v) in keys_equivalence.items():
        if k in kwds:
            kwds[v] = kwds.pop(k)
    if valid_keys is None or all(k in valid_keys for k in kwds):
        return kwds
    else:
        invalid = set(kwds) - valid_keys
        raise TypeError("Invalid keyword argument, '%s'" % invalid.pop())# Não consegui colocar qual argumento não é válido

def iter_kwds(kwds, iter_keys={}, valid_keys=None):
    '''
    Useful to create vectorized version of the keyword arguments of some 
    functions. 
    
    If iter_keys is None, all keys are considered possible iter_keys.
    
    Example
    -------
    
    Prepare the keys_equivalence and kwds_set parameters
    
    >>> iter_keys = set([ 'color', 'linewidth' ])
    >>> valid_keys = iter_keys.union(set(['disp']))

    If any key belonging to iter_keys is a sequence, it cycles over its 
    elements
    
    >>> it = iter_kwds({'color':['k', 'b'], 'disp': True}, iter_keys, valid_keys)
    >>> for i in range(3):
    ...     print it.next()
    {'color': 'k', 'disp': True}
    {'color': 'b', 'disp': True}
    {'color': 'k', 'disp': True}
    
    If any of the iter_keys is not a sequence, it is repeated in all iterations 
    
    >>> it = iter_kwds({'color':'k', 'linewidth': [1, 2, 3]}, iter_keys, valid_keys)
    >>> for i in range(3):
    ...     print it.next()
    {'color': 'k', 'linewidth': 1}
    {'color': 'k', 'linewidth': 2}
    {'color': 'k', 'linewidth': 3}
    '''

    # Create dictionary of iterators
    kwds_iter = {}
    for (k, v) in kwds.items():
        if (iter_keys is None or k in iter_keys) and isinstance(v, list):
            kwds_iter[k] = itertools.cycle(v)
        else:
            kwds_iter[k] = itertools.repeat(v)

    # Check if keys are valid
    if (valid_keys is not None) and (not all(k in valid_keys for k in kwds)):
        invalid = set(kwds) - valid_keys
        raise TypeError("Invalid keyword argument '%s'" % invalid.pop())# Não consegui colocar qual argumento não é válido

    while True:
        yield dict((k, v.next()) for (k, v) in kwds_iter.items())

def moments(sample, N, central=False):
    '''
    Return all moments of a given sample from 1 to N. 
    '''

    if N < 1:
        raise ValueError('N must be greater or equal to 1, got %s' % N)
    m_sample = np.array(sample, dtype=float)
    moments = np.zeros(N, dtype=float)

    # Compute mean and other moments
    if central:
        sample = np.asarray(sample, dtype=float) - moments[0]
    else:
        moments[0] = m_sample.mean()
    for i in range(1, N):
        m_sample *= sample
        moments[i] = m_sample.mean()

    return moments

def cumulants(sample, N):
    '''
    Return all moments of a given sample from 1 to N.
    
    Output
    ------
    
    Array containing [kappa_1, kappa_2, ..., kappa_N]
    '''

    mu = moments(sample, N, True)
    kappa = np.zeros(N, dtype=float)
    kappa[0] = mu[0]
    for i in xrange(1, N):
        kappa[i] = mu[i] - sum(binomial(i, j) * kappa[j] * mu[i - j - 1] for j in xrange(i - 1))
    kappa[0] = np.mean(sample)
    return kappa

def fat(m):
    '''
    fat(m) <==> m!
    '''
    if m >= 1:
        return reduce(op.mul, xrange(1, m + 1))
    elif m == 0:
        return 1
    else:
        raise ValueError('fatorial of a negative integer, %s' % m)

def fat2(m, n):
    '''
    Return m * (m + 1) * ... * (n - 1) * n
    '''
    if n >= m:
        return reduce(op.mul, xrange(m, n + 1))
    else:
        raise ValueError('m > n')

def binomial(n, m):
    '''
    Return the binomial coefficient
    
        binomial(n, m) = (n! / (n - m)! m !)
    '''
    if m > n:
        raise ValueError('m > n (%s > %s)' % (m, n))
    if m == 0:
        return 1
    else:
        return float(fat2(m, n)) / fat(n - m)


def time_range(t0=None, tf=None, delta_t=None, size=None):
    '''
    Return a tuple (t0, tf, delta_t, size) inferring the missing entries from 
    the given data.
    
    If necessary and not specified, it assumes that t0=0 and delta_t=1. It will
    raise a ValueError if the set of parameters are missing or overspecified. 
    '''

    t0_f = float(t0 if t0 is not None else 0.0)
    tf_f = float(tf)
    delta_tf = float(delta_t if delta_t is not None else 1.0)
    size_i = int(size)

    if size is None:
        if tf is None:
            raise ValueError('missing parameters')

        size = int((tf - t0) / delta_tf)
        return (t0_f, t0_f + size * delta_tf, delta_tf, size)
    else:
        if tf is None:
            return (t0_f, t0_f + size_i * delta_tf, delta_tf, size_i)
        elif delta_t is None:
            return (t0_f, tf_f, (tf_f - t0_f) / size_i, size_i)
        elif t0 is None:
            return (tf_f - size_i * delta_tf, tf_f, delta_tf, size_i)
        else:
            raise ValueError('overspecified parameters')

if __name__ == '__main__':
    import doctest
    doctest.testmod(optionflags=doctest.REPORT_ONLY_FIRST_FAILURE)
    from scipy.stats import kstat
#
#    S = np.random.normal(size=100000)
#    print moments(S, 10)
#    print cumulants(S, 5)
#    print [ kstat(S, i) for i in [1, 2, 3, 4]]

    t0, tf, dt, n = time_range(0, 1, size=10)
    print np.arange(t0, tf, dt)
